#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/rsa.h>
#include <openssl/err.h>

#define BUFFER_SIZE 2048

void print_openssl_error(const char *msg) {
    fprintf(stderr, "error%s\n", msg);
    ERR_print_errors_fp(stderr);
}

void generate_rsa_keypair(const char *priv_file, const char *pub_file) {
    if (access(priv_file, F_OK) == 0 && access(pub_file, F_OK) == 0) {
        printf("Bob 密钥已存在，跳过生成。\n");
        return;
    }

    printf("生成 Bob 的 RSA 密钥对...\n");
    EVP_PKEY *pkey = EVP_PKEY_new();
    RSA *rsa = RSA_new();
    BIGNUM *bn = BN_new();
    BN_set_word(bn, RSA_F4);
    RSA_generate_key_ex(rsa, 2048, bn, NULL);
    EVP_PKEY_assign_RSA(pkey, rsa);

    FILE *priv = fopen(priv_file, "wb");
    PEM_write_PrivateKey(priv, pkey, NULL, NULL, 0, NULL, NULL);
    fclose(priv);

    FILE *pub = fopen(pub_file, "wb");
    PEM_write_PUBKEY(pub, pkey);
    fclose(pub);

    EVP_PKEY_free(pkey);
    BN_free(bn);
    printf("Bob 密钥: %s, %s\n", priv_file, pub_file);
}

int main() {
    ERR_load_crypto_strings();
    generate_rsa_keypair("bob_private.key", "bob_public.key");

    // 检查是否存在 Alice 发来的信封
    if (access("message_to_bob", F_OK) != 0) {
        printf("未收到 Alice 的数字信封./alice\n");
        return 0;
    }

    FILE *in = fopen("message_to_bob", "rb");
    int iv_len, ciphertext_len;
    size_t envelope_len, signature_len;
    unsigned char aes_iv[16], ciphertext[BUFFER_SIZE], envelope[BUFFER_SIZE], signature[BUFFER_SIZE];

    fread(&iv_len, sizeof(int), 1, in);
    fread(aes_iv, 1, iv_len, in);
    fread(&ciphertext_len, sizeof(int), 1, in);
    fread(ciphertext, 1, ciphertext_len, in);
    fread(&envelope_len, sizeof(size_t), 1, in);
    fread(envelope, 1, envelope_len, in);
    fread(&signature_len, sizeof(size_t), 1, in);
    fread(signature, 1, signature_len, in);
    fclose(in);

    // 1. 用 Bob 私钥解密 AES 密钥
    FILE *bob_priv = fopen("bob_private.key", "r");
    EVP_PKEY *bob_privkey = PEM_read_PrivateKey(bob_priv, NULL, NULL, NULL);
    fclose(bob_priv);

    unsigned char aes_key[BUFFER_SIZE];
    size_t aes_key_len = sizeof(aes_key);

    EVP_PKEY_CTX *dec_ctx = EVP_PKEY_CTX_new(bob_privkey, NULL);
    EVP_PKEY_decrypt_init(dec_ctx);
    EVP_PKEY_CTX_set_rsa_padding(dec_ctx, RSA_PKCS1_PADDING);
    if (EVP_PKEY_decrypt(dec_ctx, aes_key, &aes_key_len, envelope, envelope_len) <= 0) {
        print_openssl_error("RSA 解密失败");
        return 1;
    }
    EVP_PKEY_CTX_free(dec_ctx);
    EVP_PKEY_free(bob_privkey);

    // 2. AES 解密
    unsigned char plaintext[BUFFER_SIZE];
    int len, plaintext_len;
    EVP_CIPHER_CTX *aes_ctx = EVP_CIPHER_CTX_new();
    EVP_DecryptInit_ex(aes_ctx, EVP_aes_256_cbc(), NULL, aes_key, aes_iv);
    EVP_DecryptUpdate(aes_ctx, plaintext, &len, ciphertext, ciphertext_len);
    plaintext_len = len;
    EVP_DecryptFinal_ex(aes_ctx, plaintext + len, &len);
    plaintext_len += len;
    plaintext[plaintext_len] = '\0';
    EVP_CIPHER_CTX_free(aes_ctx);

    // 3. 验签
    FILE *alice_pub_file = fopen("alice_public.key", "r");
    if (!alice_pub_file) {
        printf("缺少 alice_public.key，请从 Alice 处获取。\n");
        return 1;
    }
    EVP_PKEY *alice_pubkey = PEM_read_PUBKEY(alice_pub_file, NULL, NULL, NULL);
    fclose(alice_pub_file);

    EVP_MD_CTX *verify_ctx = EVP_MD_CTX_new();
    EVP_DigestVerifyInit(verify_ctx, NULL, EVP_sha256(), NULL, alice_pubkey);
    EVP_DigestVerifyUpdate(verify_ctx, plaintext, plaintext_len);
    int verify_ok = EVP_DigestVerifyFinal(verify_ctx, signature, signature_len);
    EVP_MD_CTX_free(verify_ctx);
    EVP_PKEY_free(alice_pubkey);

    printf("\n解密内容:\n%s\n", plaintext);
    if (verify_ok == 1)
        printf("签名验证成功：来自 Alice，内容未被篡改。\n");
    else
        printf("签名验证失败！\n");

    ERR_free_strings();
    return 0;
}
